Skip to content

[perf] Copy via memoryview in pack_into instead of torch.frombuffer - #154

Merged
0oshowero0 merged 1 commit into
Ascend:mainfrom
Chase-Rong:perf/pack-into-memoryview
Aug 12, 2026
Merged

[perf] Copy via memoryview in pack_into instead of torch.frombuffer#154
0oshowero0 merged 1 commit into
Ascend:mainfrom
Chase-Rong:perf/pack-into-memoryview

Conversation

@Chase-Rong

Copy link
Copy Markdown
Contributor

Problem

pack_into concatenates the msgpack fragments of one object into a contiguous buffer —
a 4-byte header, an 8-byte (offset, length) entry per fragment, then tightly packed
payload. It is a memcpy with an index table.

The old implementation did that memcpy through tensors:

target_tensor = torch.frombuffer(target_mv, dtype=torch.uint8)
for item in items:
    src_tensor = torch.frombuffer(memoryview(item), dtype=torch.uint8)
    target_tensor[payload_offset : payload_offset + nbytes].copy_(src_tensor)

At TQ's object sizes (one key per (sample, field), so a few hundred bytes each) none of
that work is the copy. Timed per step, 1024 objects of 512B, best of 5:

step cost what it actually does
torch.frombuffer(item_mv) 2.80us allocate TensorImpl, bind storage, buffer protocol
target_tensor[off:off+n] 3.14us another TensorImpl plus sizes/strides/storage_offset
.copy_(src) 2.13us ATen dispatch, dtype/device checks, kernel selection
total 8.07us

A 512-byte memcpy is tens of nanoseconds. Two tensor objects get constructed and one
dispatcher lookup happens per fragment, and slicing the target costs more than the copy
itself. Tensors are the wrong abstraction here — their fixed per-op cost does not
amortise over a few hundred bytes.

Change

Copy through memoryview slice assignment, which resolves to PyBuffer_ToContiguous ->
memcpy with no object allocation and no dispatcher:

step cost
memoryview(item).cast("B") 0.43us
target_mv[off:off+n] = item_mv 0.78us
total 1.20us (6.70x)

Casting both sides to "B" is required, not cosmetic: slice assignment demands matching
format and itemsize, and callers pass bytes, a memoryview, or a float32 numpy array.
Normalising to unsigned bytes makes the assignment legal and correct for all three.

Second change in the same function: the old code built a memoryview per fragment
twice — once inside calc_packed_size just to read .nbytes, once in the copy loop.
The views are now materialised once and shared, with required computed inline.

Measurements

verl GRPO on main at 750719e, this PR applied alone (the other two perf branches
reverted to baseline). Medians; two independent baseline runs per configuration are
shown to give the run-to-run spread.

Single-node, 512-sample put (1024 objects):

baseline this PR
pack_into, per object 50.4us 8.2us 6.1x
calc_packed_size calls 2048 1024 halved
calc_packed_size total 8.6ms 2.4ms 3.58x
pack phase wall clock 188.9 / 192.5ms 35.5ms 5.37x
TQ-side total 269.5 / 275.3ms 113.2ms 2.41x
put_data 316 / 319ms 163ms 1.95x

Dual-node, 1024-sample put (2048 objects):

baseline this PR
pack phase wall clock 364.5 / 367.9ms 97.8ms 3.74x
TQ-side total 530.8 / 529.9ms 245.1ms 2.16x

The per-step micro-benchmark ratio (6.70x) and the production ratio (6.1x) agree, so the
mechanism above accounts for the result. Absolute values differ because production
fragments are larger and more numerous than the synthetic 512B case.

yuanrong RPC time is unchanged (32.1 -> 31.5ms single-node, 55.2 -> 61.6ms dual-node),
so the saving is entirely client-side CPU. get_data is unaffected (medians within 3%
across variants).

Note dual-node put_data wall clock is not a reliable metric in this environment: an
intermittent ~480ms run_in_executor handoff delay appears in some runs (two baseline
runs measured 606ms and 1083ms with every instrumented step within 1%, the difference
landing entirely in un-instrumented time). Judge the dual-node effect by TQ-side total
or pack phase, which agree to 0.2%/0.9% between those two runs.

Side effect: GIL contention drops as well

batch_encode_into still dispatches pack_into across 16 threads. Because each task now
holds the GIL for 1.20us instead of 8.07us, the queue stops backing up:

real work 16-thread pack_into sum (incl. lock wait) pack phase wall clock
baseline 51.6ms 1988.5ms 188.9ms
this PR ~6.6ms 6.6ms 35.5ms

The baseline's per-call wall clock inflates to 1942us under contention; here the
accumulated per-thread time equals the real work, i.e. threads are barely waiting on each
other. That is why the phase improves by more than the 6.7x of the copy alone.

Correctness

Output is byte-for-byte identical to the previous implementation. Both
tests/test_serial_utils_on_cpu.py and tests/test_serial_utils_batch_on_cpu.py pass,
and packed buffers still round-trip through unpack_from/decode unchanged. The
buffer-too-small `ValueError

pack_into built a uint8 tensor per item with torch.frombuffer and dispatched a
copy_ kernel for it. At TQ's object sizes the dispatch dominates: one key is
stored per (field, sample), so items are a few hundred bytes each and the tensor
construction plus kernel launch costs far more than the memcpy.

Copy through memoryview slice assignment instead. Casting both sides to "B"
keeps this correct whether the caller passes bytes, a memoryview or a numpy
array. Also compute the size requirement from the item views already
materialised for the copy, so calc_packed_size no longer walks every item a
second time.

Measured (Ascend NPU container, torch 2.9.0+cpu):

   1024 objects x 512B:   35.0ms -> 5.5ms  (6.4x)
   7168 objects x 512B:  248.6ms -> 40.7ms (6.1x)
   2048 objects x 16KB:   82.5ms -> 22.5ms (3.7x)

In a real verl GRPO run, pack_into for a 512-sample 2-field put went 51.3ms ->
6.0ms, taking put_data from 183.9ms to 139.7ms single-node and 364.9ms to
290.4ms dual-node.

Output is byte-for-byte identical to the previous implementation; both
serial_utils test modules pass, and the packed buffers still round-trip through
unpack_from/decode unchanged.

Note both callers benefit: YuanrongStorageClient and MooncakeStore both go
through batch_encode_into. The measurements above are from the yuanrong path;
Mooncake was not measured.
@0oshowero0
0oshowero0 requested a balanced review from Copilot August 12, 2026 08:20

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@0oshowero0

Copy link
Copy Markdown
Collaborator

@codex

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: bb041f6b1a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@0oshowero0
0oshowero0 merged commit 2fc5b79 into Ascend:main Aug 12, 2026
7 of 8 checks passed
@ascend-robot

Copy link
Copy Markdown

CLA Signature Guide

@Chase-Rong , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit Reason
[bb041f6 [perf] Copy via memoryview in p...](bb041f6) the email used in the commit is not linked to a signed CLA!
please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants